import os
import gradio as gr
from openai import OpenAI
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize OpenAI client
client = OpenAI(
    api_key=os.getenv("LLM_API_KEY"),
    base_url=os.getenv("LLM_BASE_URL"),
)

def generate_joke(theme, style, length):
    if not theme.strip():
        return "请先输入一个主题。"
    length_desc = ["短", "中", "长"][length-1]
    prompt = f"Generate a {length_desc} funny joke about {theme} in {style} style. Keep it clean."
    try:
        response = client.chat.completions.create(
            model=os.getenv("LLM_MODEL"),
            messages=[{"role": "user", "content": prompt}],
            max_tokens=5000,  # Increased for longer jokes
        )
        joke = response.choices[0].message.content.strip()
        return joke, gr.update(interactive=True)
    except Exception as e:
        return f"Error: {str(e)}", gr.update(interactive=False)

def regenerate_joke(theme, style, length):
    return generate_joke(theme, style, length)

# Gradio Blocks interface for better control
with gr.Blocks(title="AI 笑话生成器", theme=gr.themes.Soft()) as iface:
    gr.HTML("<h2 style='text-align: center; color: #4CAF50;'>🤖 AI 笑话生成器</h2>")
    gr.HTML("<p style='text-align: center;'>输入主题、风格和长度，生成有趣的笑话。</p>")

    with gr.Row():
        theme_input = gr.Textbox(label="主题 (Theme)", placeholder="输入主题，如动物、工作等")
        style_input = gr.Dropdown(
            label="风格 (Style)",
            choices=["幽默的 (Humorous)", "讽刺的 (Sarcastic)", "搞笑的 (Funny)", "轻松的 (Light-hearted)", "双关语 (Pun)", "夸张的 (Exaggerated)", "随意的 (Casual)"],
            value="幽默的 (Humorous)"
        )
        length_input = gr.Slider(label="长度 (Length)", minimum=1, maximum=3, step=1, value=2, info="1=短, 2=中, 3=长")

    generate_btn = gr.Button("生成笑话", variant="primary")

    joke_output = gr.Textbox(label="生成的笑话 (Generated Joke)", lines=5, autoscroll=False, interactive=False)

    regen_btn = gr.Button("🔄 重新生成 (Regenerate)", interactive=False)

    # Event handlers
    generate_btn.click(
        fn=generate_joke,
        inputs=[theme_input, style_input, length_input],
        outputs=[joke_output, regen_btn],
        show_progress=True
    )

    regen_btn.click(
        fn=regenerate_joke,
        inputs=[theme_input, style_input, length_input],
        outputs=[joke_output, regen_btn]
    )

if __name__ == "__main__":
    iface.launch()
